import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.conv import _ConvNd
from torch.nn.modules.utils import _pair
from torchvision.models import resnet34

from networks.Fuck import ResNet18WithoutMaxPool, ResNet34WithoutMaxPool


def base_upsample(x, size):
    return F.interpolate(x, size=size, mode='bilinear', align_corners=True)


class BaseConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1):
        super(BaseConvBNReLU, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels,
                      out_channels,
                      kernel_size=kernel_size,
                      stride=stride,
                      padding=padding,
                      dilation=dilation,
                      bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class GatedSpatialConv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=False):
        super(GatedSpatialConv2d, self).__init__(in_channels, out_channels,
                                                 _pair(kernel_size),
                                                 _pair(stride),
                                                 _pair(padding),
                                                 _pair(dilation),
                                                 False, _pair(0), groups, bias, 'zeros')

        self.gate_conv = nn.Sequential(
            nn.BatchNorm2d(in_channels + 1),
            nn.Conv2d(in_channels + 1, in_channels + 1, 1),
            nn.ReLU(True),
            nn.Conv2d(in_channels + 1, 1, 1),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

    def forward(self, input_features, gating_features):
        """
        :param input_features:  [NxCxHxW]
        :param gating_features: [Nx1xHxW]
        :return:
        """
        alphas = self.gate_conv(torch.cat([input_features, gating_features], dim=1))
        input_features = (input_features * (alphas + 1))
        return F.conv2d(input_features, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)

    def reset_parameters(self):
        nn.init.xavier_normal_(self.weight)
        if self.bias is not None:
            nn.init.zeros_(self.bias)


class ShitASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ShitASPP, self).__init__()

        rate1, rate2, rate3 = (6, 12, 18)

        self.conv_list = nn.ModuleList([
            BaseConvBNReLU(in_channels, out_channels, 1),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate1, rate1),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate2, rate2),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate3, rate3),
            nn.Sequential(
                BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, 1),
                BaseConvBNReLU(out_channels // 2, out_channels // 2, 5, 1, 2)
            ),
        ])
        self.pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            BaseConvBNReLU(in_channels, out_channels * 3, 1),
        )

        self.project = BaseConvBNReLU(3 * out_channels, out_channels, 1)

    def forward(self, x):
        size = x.size()[2:]
        res = []
        for conv in self.conv_list:
            res.append(F.interpolate(conv(x), size=size, mode='bilinear', align_corners=False))
        res = torch.cat(res, dim=1)
        pool = base_upsample(self.pool(x), size)
        return self.project(res + pool)


class DetailExtraction(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DetailExtraction, self).__init__()
        self.conv1 = BaseConvBNReLU(in_channels, out_channels, 3, 1, 1)
        self.conv2 = BaseConvBNReLU(in_channels, out_channels, 1, 1, 0)
        self.conv3 = BaseConvBNReLU(out_channels * 2, out_channels, 3, 1, 1)

    def forward(self, x):
        f1 = self.conv1(x)
        f2 = self.conv2(x)
        out = self.conv3(torch.cat([f1, f2], dim=1))
        return out


class ShitMiddle(nn.Module):
    def __init__(self, x1_channels, x4_channels, fine_channels):
        super(ShitMiddle, self).__init__()
        self.aspp = ShitASPP(x4_channels, x4_channels)
        self.detail_extraction = DetailExtraction(x1_channels, fine_channels)
        self.down = BaseConvBNReLU(x4_channels, fine_channels, 1)
        self.edge = BaseConvBNReLU(fine_channels * 2, 1, 1)

    def forward(self, x1, x4):
        f1 = self.detail_extraction(x1)
        f4 = self.aspp(x4)

        down = self.down(f4)
        edge = self.edge(torch.cat([base_upsample(f1, f4.size()[2:]), down], dim=1))
        return [f1, f4, edge]


class SE(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SE, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)


class ShitDecoder(nn.Module):
    def __init__(self, in_channels, fine_channels, num_classes):
        super(ShitDecoder, self).__init__()
        self.de_en_conv = nn.Sequential(
            BaseConvBNReLU(in_channels * 3, in_channels * 3, 3, 1, 1),
            SE(in_channels * 3),
            BaseConvBNReLU(in_channels * 3, in_channels, 3, 1, 1),
        )
        self.edge = BaseConvBNReLU(in_channels, 1, 1)
        self.down = BaseConvBNReLU(in_channels, fine_channels, 3, 1, 1)
        self.new_out = nn.Sequential(
            BaseConvBNReLU(in_channels + fine_channels, in_channels, 3, 1, 1),
            SE(in_channels),
            BaseConvBNReLU(in_channels, in_channels, 3, 1, 1),
        )
        self.gate = GatedSpatialConv2d(fine_channels, fine_channels, 3, 1, 1)
        self.sup = BaseConvBNReLU(in_channels, num_classes, 1)

    def forward(self, en, de, fine):
        de = base_upsample(de, en.size()[2:])
        de_en = self.de_en_conv(torch.cat([de, en], dim=1))
        edge = base_upsample(self.edge(de_en), fine.size()[2:])
        down = fine * base_upsample(self.down(de_en), fine.size()[2:])
        new_out = self.new_out(torch.cat([de_en, base_upsample(down, de_en.size()[2:])], dim=1))
        new_fine = self.gate(down + fine, edge)
        sup = self.sup(new_out)
        return [new_fine, new_out, sup]


class ShitRes34(nn.Module):
    def __init__(self, num_classes, pretrained=None):
        super(ShitRes34, self).__init__()
        self.backbone = ResNet34WithoutMaxPool(pretrained=pretrained)
        in_channels = self.backbone.in_channels
        fine_channels = 128
        self.middle = ShitMiddle(in_channels[0], in_channels[3], fine_channels)
        self.decoder3 = ShitDecoder(in_channels[2], fine_channels, num_classes)
        self.decoder2 = ShitDecoder(in_channels[1], fine_channels, num_classes)
        self.decoder1 = ShitDecoder(in_channels[0], fine_channels, num_classes)
        self.edge_sup = BaseConvBNReLU(fine_channels, 1, 1)
        self.out = nn.Sequential(
            BaseConvBNReLU(in_channels[0] + fine_channels, in_channels[0], 3, 1, 1),
            nn.Conv2d(in_channels[0], num_classes, 1)
        )

    def forward(self, xx):
        x1, x2, x3, x4 = self.backbone(xx)
        fine, up, edge1 = self.middle(x1, x4)
        mid_fine = fine
        mid_up = up
        fine, up, sup3 = self.decoder3(x3, up, fine)
        fine, up, sup2 = self.decoder2(x2, up, fine)
        fine, up, sup1 = self.decoder1(x1, up, fine)
        edge2 = self.edge_sup(fine)
        out = self.out(torch.cat([fine, up], dim=1))
        out = base_upsample(out, xx.size()[2:])
        return [out, [sup1, sup2, sup3], [edge1, edge2], [mid_fine, mid_up]]


class ShitDecoderSmall(nn.Module):
    def __init__(self, in_channels, fine_channels, num_classes):
        super(ShitDecoderSmall, self).__init__()
        self.de_en_conv = nn.Sequential(
            BaseConvBNReLU(in_channels * 3, in_channels * 3, 3, 1, 1),
            SE(in_channels * 3),
            BaseConvBNReLU(in_channels * 3, in_channels, 3, 1, 1),
        )
        self.edge = BaseConvBNReLU(in_channels, 1, 1)
        self.down = BaseConvBNReLU(in_channels, fine_channels, 3, 1, 1)
        self.new_out = nn.Sequential(
            BaseConvBNReLU(in_channels + fine_channels, in_channels, 3, 1, 1),
            SE(in_channels),
            BaseConvBNReLU(in_channels, in_channels, 3, 1, 1),
        )
        self.gate = GatedSpatialConv2d(fine_channels, fine_channels, 3, 1, 1)
        self.sup = BaseConvBNReLU(in_channels, num_classes, 1)

    def forward(self, en, de, fine):
        de = base_upsample(de, en.size()[2:])
        de_en = self.de_en_conv(torch.cat([de, en], dim=1))
        edge = base_upsample(self.edge(de_en), fine.size()[2:])
        down = fine * base_upsample(self.down(de_en), fine.size()[2:])
        new_out = self.new_out(torch.cat([de_en, base_upsample(down, de_en.size()[2:])], dim=1))
        new_fine = self.gate(down + fine, edge)
        sup = self.sup(new_out)
        return [new_fine, new_out, sup]


class ShitRes18(nn.Module):
    def __init__(self, num_classes, pretrained=None):
        super(ShitRes18, self).__init__()
        self.backbone = ResNet18WithoutMaxPool(pretrained=pretrained)
        in_channels = self.backbone.in_channels
        fine_channels = 128
        self.middle = ShitMiddle(in_channels[0], in_channels[3], fine_channels)
        self.decoder3 = ShitDecoder(in_channels[2], fine_channels, num_classes)
        self.decoder2 = ShitDecoder(in_channels[1], fine_channels, num_classes)
        self.decoder1 = ShitDecoder(in_channels[0], fine_channels, num_classes)
        self.edge_sup = BaseConvBNReLU(fine_channels, 1, 1)
        self.out = nn.Sequential(
            BaseConvBNReLU(in_channels[0] + fine_channels, in_channels[0], 3, 1, 1),
            nn.Conv2d(in_channels[0], num_classes, 1)
        )

    def forward(self, xx):
        x1, x2, x3, x4 = self.backbone(xx)
        fine, up, edge1 = self.middle(x1, x4)
        mid_fine = fine
        mid_up = up
        fine, up, sup3 = self.decoder3(x3, up, fine)
        fine, up, sup2 = self.decoder2(x2, up, fine)
        fine, up, sup1 = self.decoder1(x1, up, fine)
        edge2 = self.edge_sup(fine)
        out = self.out(torch.cat([fine, up], dim=1))
        out = base_upsample(out, xx.size()[2:])
        return [out, [sup1, sup2, sup3], [edge1, edge2], [mid_fine, mid_up]]
